/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2019-2022. All rights reserved.
 * Description: all-reduce fusion header file
 * Author: lilianlin
 * Create: 2019-12-8
 */

#ifndef AUTO_TUNING_HCOM_ALL_REDUCE_FUSION_H
#define AUTO_TUNING_HCOM_ALL_REDUCE_FUSION_H

#include "hccl/hccl_types.h"
#include "common/optimizer/graph_optimizer.h"
#include "common/optimizer/graph_optimizer_types.h"
#include "graph/compute_graph.h"
#include "op_fusion_base_pub.h"
#include "../hcom/hcom_all_reduce_fusion.h"
#include "platform/platform_info.h"
#include "nlohmann/json.hpp"

namespace hccl {
using GradientDataInfo = struct GradientDataInfoDef {
    int64_t dataSize;
    std::string dataType;
    uint32_t graphId;
    std::string groupName;
    std::string gradientNodeName;
    std::string traceNodeName;
    std::string allReduceNodeName;
};

class AutoTuningHcomAllReduceFusion : public HcomAllReduceFusion {
public:
    AutoTuningHcomAllReduceFusion();
    ~AutoTuningHcomAllReduceFusion() override;
    HcclResult Run(ge::ComputeGraph& graph, std::vector<GradientDataInfo> &recordInfos);
protected:
    HcclResult FuseOps(ge::ComputeGraph& graph, FusionSection& fusionSection) override;
    HcclResult GetGradSplitStrategy(const std::string& modelName, const std::string& sGroup, \
        const FusionSection& fusionSection, u32& segmentNum, std::vector<u32>& segmentIndex) override;\
private:
    HcclResult RecordGradientDataInfo(ge::ComputeGraph& graph, std::vector<GradientDataInfo> &recordInfos);
    HcclResult GetGroupName(const ge::OpDescPtr &op, std::string &group);
    HcclResult GetDataTypeName(const ge::DataType dataType, std::string &dataTypeName);
    HcclResult GetGradientDataInfo(ge::ComputeGraph &graph, ge::NodePtr &nodePtr, GradientDataInfo &gradientNodeInfo);
    HcclResult AddTraceNode(ge::ComputeGraph &graph, ge::NodePtr &nodePtr, const GradientDataInfo gradientNodeInfo);
    HcclResult SetGradientInformation(ge::ComputeGraph &graph);
    HcclResult AddFusionMapInFusionJson(const std::string& fusionHash);
    HcclResult SetFusionModelInLibrary(std::string& workPath, const std::string& fusionHash);
    HcclResult GetFusionWorkPath(std::string &fusionPath);
    HcclResult CreateFusionDir(std::string &dir);
    bool AutoTuneTargetNode(ge::NodePtr nodePtr);
    bool isNotFoundHash_;
};
}
#endif
